#!/usr/bin/python
'''
Takes two fasta files and makes the following cases.
1: all type 1.
2: joined reads.
3: all type 2.
'''
import gzip
import sys
import os
import random

def read_file(filepath):
	# define data structure.
	data = {}
	
	# Check if zipped.
	zipped = False
	if filepath.count(".gz"):
		zipped = True
	
	# Open file.
	seq = ""
	seqheader = ""
	first = True
	
	if zipped == True:
		fin = gzip.open(filepath,"r")
	else:
		fin = open(filepath,"r")
	for line in fin:
		# Skip comments.	
		if line[0] == "#": continue

		# first time.
		if first == True and line[0] == ">":
			seq = ""
			seqheader = line.strip().replace(">","")
			first = False
			continue		

		# finished seq.
		if line[0] == ">":
			# save data.
			data[seqheader] = seq.strip()

			# Reset variables.
			seq = ""
			seqheader = line.strip().replace(">","")
			continue

		# get seq.
		seq += line.strip()
			
	# Handle last close and return.
	data[seqheader] = seq.strip()
	fin.close()
	return data
		
	
	

# Parameters.
type1_file = sys.argv[1]
type2_file = sys.argv[2]

# Read in both files.
type1_seqs = read_file(type1_file)
type2_seqs = read_file(type2_file)

# Choose smallest
sz1 = len(type1_seqs)
sz2 = len(type2_seqs)
if sz1 < sz2:
	max_sz = sz1
if sz2 <= sz1:
	max_sz = sz2

# Divide into 2.
half_sz = int( float(max_sz) / 2.0 )

# Choose random merged samples.
type1_names = type1_seqs.keys()
type2_names = type2_seqs.keys()

plain1 = random.sample(type1_names, half_sz)
plain2 = random.sample(type2_names, half_sz)

# Get names of ones to merge.
used1 = set(plain1)
used2 = set(plain2)
tot1 = set(type1_names)
tot2 = set(type2_names)

junct1 = list( tot1.difference(used1) )
junct2 = list( tot2.difference(used2) )

# Shrink each set to smallest size.
sz1 = len(plain1)
sz2 = len(plain2)
sz = min(sz1,sz2)
plain1 = plain1[0:sz]
plain2 = plain2[0:sz]

sz1 = len(junct1)
sz2 = len(junct2)
sz = min(sz1,sz2)
junct1 = junct1[0:sz]
junct2 = junct2[0:sz]

# Write out junction reads.
fout = open("junct.fasta","w")
for i in range(len(junct1)):
	# Get data.
	name1 = junct1[i]
	name2 = junct2[i]
	seq1 = type1_seqs[name1]
	seq2 = type2_seqs[name2]
	
	if len(seq1) < 100 or len(seq2) < 100:
		continue
	
	# Randomize order.
	dat = [seq1, seq2]
	random.shuffle(dat)
	
	# Combine reads at one point.
	bp = random.randint(20, len(dat[0]) - 20)
	jseq = dat[0][0:bp]
	jseq += dat[1][bp:len(dat[1])-1]
	
	# Write out new read.
	fout.write(">%s_%s\n%s\n" % (name1, name2, jseq) )
	
fout.close()

# Print out plain.
fout = open("plain_1.fasta", "w")
for name in plain1:
	fout.write(">%s\n%s\n" % (name, type1_seqs[name]) )
fout.close()

fout = open("plain_2.fasta", "w")
for name in plain2:
	fout.write(">%s\n%s\n" % (name, type2_seqs[name]) )
fout.close()


	
